import sys
import os

# 获取项目根目录的路径（即包含 scripts 和 modules 的目录）
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(project_root)

# 本文件的目的是用来测试加载网络看是否能加载成功
from modules.activation import get_activation
from modules.actor_critic import *
from modules.depth_backbone import *
from modules.estimator import *
import argparse
import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

def test(args):
    depth_backbone = DepthOnlyFCBackbone58x87(32)
    depth_model = RecurrentDepthBackbone(depth_backbone, args.n_proprio)
    actor_model = Actor(args.n_proprio, args.n_scan, args.n_actions,  [128, 64, 32],
                        [512, 256, 128], [64, 20], args.n_priv_latent,
                        args.n_priv_explicit, args.n_hist, get_activation("elu"))
    estimator_model = Estimator(args.n_proprio, args.n_priv_explicit, [128, 64], "elu")
    load_path = "/home/mi/MyRL/models/model_17000.pt"
    as_state_dict = torch.load(load_path)

    depth_model.to(device)
    actor_model.to(device)
    estimator_model.to(device)
    depth_model.load_state_dict(as_state_dict['depth_encoder_state_dict'])
    actor_model.load_state_dict((as_state_dict['depth_actor_state_dict']))
    estimator_model.load_state_dict(as_state_dict['estimator_state_dict'])
    obs_input = torch.ones(1, args.n_obs, device=device)
    depth_image = torch.ones(1, 58, 87, device=device)
    obs_prop = obs_input[:, :args.n_proprio].clone()
    obs_prop[:, 6:8] = 0
    depth_latent_and_yaw = depth_model(depth_image, obs_prop)
    depth_latent = depth_latent_and_yaw[:, :-2]
    yaw = depth_latent_and_yaw[:, -2:]
    obs_input[:, 6:8] = 1.5 * yaw

    obs_priv = estimator_model(obs_input[:, :args.n_proprio])
    # print(obs_priv.shape)
    # print(obs_input[:, args.n_proprio + args.n_scan:args.n_proprio + args.n_scan + args.n_priv_explicit].shape)
    # print(args.n_proprio + args.n_scan)
    # print(args.n_proprio + args.n_scan + args.n_priv_explicit)
    obs_input[:, args.n_proprio + args.n_scan:args.n_proprio + args.n_scan + args.n_priv_explicit] = obs_priv

    while True:
        obs_output = actor_model(obs_input, True, False, depth_latent)
        print(obs_output.shape)
        print(obs_output)





if __name__ =="__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--n_obs', type=int, default=753)
    parser.add_argument('--n_proprio', type=int, default=53)
    parser.add_argument('--n_scan', type=int, default=132)
    parser.add_argument('--n_actions', type=int, default=12)
    parser.add_argument('--n_priv_latent', type=int, default=29)
    parser.add_argument('--n_priv_explicit', type=int, default=9)
    parser.add_argument('--n_hist', type=int, default=10)
    args = parser.parse_args()
    test(args)